from .SSP import SSP_MatchingNet
from .PANet import PANet
from .PFENet import PFENet
from .CPMT import CPMT

def build_model(args):
    if args.model == 'ssp':
        model=SSP_MatchingNet(backbone=args.backbone, pretrained=args.pretrained,refine=args.refine)
    elif args.model == 'panet':
        model=PANet(backbone=args.backbone, shot=args.shot, pretrained=args.pretrained)
    elif args.model=="pfenet":
        model=PFENet(backbone=args.backbone, shot=args.shot, pretrained=args.pretrained)
    elif args.model=="cpmt":
        model=CPMT(backbone=args.backbone, shot=args.shot, pretrained=args.pretrained)
    else:
        raise NotImplementedError
     
    return model





